import numpy as np
import torchvision.datasets as datasets
import torch
import torch.nn as nn
from torchvision import transforms
from torch.utils.data.dataloader import DataLoader
from model.Vit import *
import torch.optim as optim
import matplotlib.pylab as plt
plt.switch_backend('agg')
classes = ('plane','car','bird','cat','deer',
          'dog','frog','horse','ship','truck')

 # 设置transforms
transform = transforms.Compose([
    transforms.ToTensor(), # numpy -> Tensor
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))  # 归一化 ，范围[-1,1]
])
# 训练集
trainset = datasets.CIFAR10(root='./CIFAR10',train=True,download=False,transform=transform)
# 测试集
# testset = datasets.CIFAR10(root='./CIFAR10',train=False,download=False,transform=transform)
BATCH_SIZE = 2

train_loader = DataLoader(trainset,batch_size=BATCH_SIZE,shuffle=True,num_workers=0)

# test_loader = DataLoader(testset,batch_size=BATCH_SIZE,shuffle=True,num_workers=4)

model=Vision_Model()
model=model.train()
# model=model.cuda()
optimizer=optim.Adam(model.parameters(),lr=0.0001)
Loss=nn.CrossEntropyLoss()
list1=[]
list2=[]
print(len(trainset))
if __name__=='__main__':

    for epoch in range(10):
        sumloss=0
        sum1=0
        print(epoch)
        for i,(x,y) in enumerate(train_loader):
            # x=x.cuda()
            # y=y.cuda()
            y_pred=model(x)
            # y_pred=y_pred[:,0]
            optimizer.zero_grad()
            yy=torch.argmax(y_pred,dim=1)
            num=sum(yy==y)
            sum1+=num
            loss=Loss(y_pred,y)
            sumloss+=loss.item()
            loss.backward()
            optimizer.step()
        acc=sum1/50000.0
        list1.append(sumloss/25000.0)
        list2.append(acc)
        print(acc,sumloss/25000.0)
        plt.plot(list1)
        plt.savefig('loss.jpg')
        plt.cla()
        plt.plot(list2)
        plt.savefig('acc.jpg')